# author：马楚云
import re
import os
import json
import sqlite3
import psutil
import subprocess
from configobj import ConfigObj
from .Comparator import Comparator

class BaseLineCheck(object):
    # 初始化
    def __init__(self):
        self.__config_path = "./resources/secpolicy.cfg"    # 配置文件路径
        self.__sqlite_db_path = "./resources/configuration.db"  # sqlite中基线检查规则的路径
        self.__check_criterion = self.__get_check_criterion()
        self.__comparator = Comparator()    # 比较类
        self.__check_res = None     # 检测结果
        self.__security_code = None # 安全级别

    # 连接sqlite数据库
    def __connect_sqlite(self):
        connection = sqlite3.connect(self.__sqlite_db_path)
        cursor = connection.cursor()
        return connection, cursor

    # 取出基线检查标准
    def __get_check_criterion(self):
        connection, cursor = self.__connect_sqlite()
        cursor.execute("SELECT * FROM configuration")
        results = cursor.fetchall()
        cursor.close()
        connection.close()
        return results

    # 生成本机配置文件
    def __get_cfg_file(self):
        utf16_file_path = self.__config_path
        cmd = f"secedit /export /cfg {utf16_file_path}"
        subprocess.Popen(cmd, shell=True).wait()
        self.__utf16le_to_utf8(utf16_file_path, self.__config_path)

    @staticmethod
    # 读取utf16编码文件，输出utf8编码文件
    def __utf16le_to_utf8(input_path, output_path):
        try:
            with open(input_path, "r", encoding='utf-16 le') as infile:
                content = infile.read()
            with open(output_path, "w", encoding='utf-8') as outfile:
                content = content.replace(",,", ",")
                outfile.write(content)
        except UnicodeDecodeError:
            pass
    
    # 读取配置文件
    def __cfg_read(self, section, key):
        conf_obj = ConfigObj(self.__config_path, encoding="utf-8")
        try:
            return conf_obj[section][key]
        except KeyError as e:
            return 'Key not exist'

    @staticmethod
    # 获取主机提供服务的配置信息
    def __get_service_config(cmd):
        # 1060 error code
        res = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                               stderr=subprocess.STDOUT, shell=False)
        config_val = res.communicate()[0].decode('gbk')
        # 判断服务是否存在
        if "SERVICE_NAME" not in config_val and "1060" in config_val:
            ret_config_val = None
            ret_config_val_name = "服务未安装！"
        # 获得服务信息
        else:
            val_pattern = "\\s*START_TYPE\\s*:\\s*(\\d)"
            name_patter = "\\s*START_TYPE\\s*:\\s*\\d\\s*(\\S*)"
            ret_config_val = re.findall(val_pattern, config_val)[0]
            ret_config_val_name = re.findall(name_patter, config_val)[0]
        return ret_config_val, ret_config_val_name

    # 通过sc命令获取已安装服务的信息
    def __check_services_by_sc(self):
        check_res = []
        check_criterion_list = [
            row for row in self.__check_criterion if row[-2] == 'Services']
        # 遍历已安装的服务的名单
        for row in check_criterion_list:
            proposed_val = row[6]
            operator = row[7]
            cmd = f'sc qc {row[-1]}'  # 检查该服务信息的shell命令
            config_val, config_val_name = self.__get_service_config(cmd)
            # 获取核查结果
            comp_res = self.__comparator.compare(
                config_val, proposed_val, operator)
            check_res.append([row[0], row[1], config_val_name,
                             row[5], comp_res, row[9], row[10]])
        check_res.sort(key=lambda row: row[0])
        return check_res

    # 基线安全策略检测
    def __check_sec_policy(self):
        check_res = {
            'System Access': [],
            'Event Audit': [],
            'Registry Values': [],
            'Privilege Rights': []
        }
        check_criterion_list = [row for row in self.__check_criterion if
                                row[-2] in ['System Access', 'Registry Values', 'Privilege Rights']]
        for row in check_criterion_list:
            # 遍历整个检查条目信息列表
            section = row[11]  # 检测的节
            key = row[12]  # 检测的key
            proposed_val = row[6]  # 核查条目中建议值
            operator = row[7]  # 该条目的比较操作符
            policy_name = row[1]  # 策略名称
            proposed_comment = row[5]  # 策略建议说明
            config_val = self.__cfg_read(section, key)  # 从配置文件当中获取到的配置value
            if config_val != 'Key not exist':
                # 调用comparator的compare方法，传入检测值以及建议值和比较操作符，返回检测结果
                comp_res = "符合" if config_val is None or self.__comparator.compare(
                    config_val, proposed_val, operator) else "不符合"
            else:
                config_val = "未检测到配置信息"
                comp_res = "不符合"
                
            # se比较符处理较为特殊，单独处理检测结果
            if operator == 'se':
                config_val = list(
                    set(self.__comparator.get_config_roles(config_val)))
                config_val_str = ''
                for val in config_val:
                    config_val_str += val.capitalize() + ','
                config_val = config_val_str[:-1]
                proposed_comment = row[6]
            # 将检测的结果插入到对应核查条目中
            check_res[section].append(
                [row[0], policy_name, config_val, proposed_comment, comp_res, row[9], row[10]])
        return check_res

    # 基线检查
    def check(self):
        self.__get_cfg_file()
        check_res = self.__check_sec_policy()
        check_res['Services'] = self.__check_services_by_sc()
        bc_result = {key: 0 for key in check_res.keys()}
        is_according = 1    # 是否检测通过
        according_count = 0     
        for key in check_res.keys():
            for res in check_res[key]:
                # 统计符合基线检查查条件的条数
                if res[4] == '符合':
                    bc_result[key] += 1
                else:
                    is_according = 0
            according_count += bc_result[key]   # 检测通过数量
        # 判断核查结果
        if according_count < 50:
            # 低风险
            self.__security_code = 0
        elif 50 <= according_count < 150:
            # 中风险
            self.__security_code = 1
        else:
            # 高风险
            self.__security_code = 2
        bc_result['is_according'] = is_according    # 是否检查
        self.__check_res = bc_result
        
        # 对返回的json数据重新进行格式化处理，方便前端数据处理
        return_check_res = []
        index = 1
        for key in check_res.keys():
            for res in check_res[key]:
                res.append(key)
                res.append(index)
                index += 1
                return_check_res.append(res)
        return json.dumps(return_check_res, ensure_ascii=False)

    # 获取基线检查统计结果，返回json字符串
    def get_bc_result(self):
        meta_data = {
            "services": self.__check_res['Services'],
            "systemAccess": self.__check_res['System Access'],
            "privilegeRights": self.__check_res['Privilege Rights'],
            "registryValue": self.__check_res['Registry Values'],
            "isAccording": self.__check_res['is_according'],
            "securityCode": self.__security_code
        }
        return json.dumps(meta_data)

# 主函数
def main():
    # 创建一个基线检查类
    base_line_checker = BaseLineCheck()
    # 获得基线检查结果
    check_res = base_line_checker.check()
    # 获得基线检查统计结果
    result = base_line_checker.get_bc_result()
    # 打印结果
    print(check_res)
    print(result)

if __name__ == '__main__':
    main()
